import os
import jieba
import pandas as pd
from collections import Counter
import matplotlib.pyplot as plt
from pylab import *
mpl.rcParams['font.sans-serif'] = ['SimHei']

# 加载停用词表
stopwords = []
with open('thesaurus/stop_word.txt', 'r', encoding='utf-8') as f:
    for line in f:
        stopwords.append(line.strip())

# 加载自定义词典
jieba.load_userdict('thesaurus/user_word.txt')
# 读取xlsx文件
df = pd.concat([pd.read_excel('document/' + file_name) for file_name in os.listdir('document/') if file_name.endswith('.xlsx')])
# 遍历document文件夹中的xlsx文件
def remove_float(x):
    if isinstance(x, float):
        return ''
    else:
        return x

def cut_words(text):
    words = jieba.cut(text)
    return [word for word in words if word not in stopwords and len(word) > 1]

# 定义计算主题强度函数
def calculate_topic_strength(df, topic):
    # 将content列文本用jieba进行分词并删除停用词
    df['content'] = df['content'].apply(remove_float)
    df['words'] = df['content'].apply(cut_words)

    # 计算主题词出现的次数
    topic_count = Counter([word for words in df['words'] for word in words if word == topic])
    # 计算总词数
    total_count = sum([len(words) for words in df['words']])
    # 计算主题强度
    topic_strength = topic_count[topic] / total_count
    return topic_strength

# 计算主题强度并绘制可视化图表、生成csv文件和词云图
for i, filename in enumerate(os.listdir('document')):
    if filename.endswith('.xlsx'):
        # 获取新文件名
        new_title = f'{filename[:-5]}_{i}'
        # 读取xlsx文件
        df = pd.read_excel(f'document/{filename}')
        # 计算主题强度
        economic_strength = calculate_topic_strength(df, '经济')
        livelihood_strength = calculate_topic_strength(df, '民生')
        education_strength = calculate_topic_strength(df, '教育')
        society_strength = calculate_topic_strength(df, '社会')
        military_strength = calculate_topic_strength(df, '军事')
        culture_strength = calculate_topic_strength(df, '文化')
        technology_strength = calculate_topic_strength(df, '科技')
        politics_strength = calculate_topic_strength(df, '政治')

        # 绘制可视化图表
        labels = ['经济', '民生', '教育','社会','军事','文化','科技','政治']
        values = [economic_strength, livelihood_strength, education_strength,society_strength,military_strength , culture_strength ,technology_strength,  politics_strength]
        plt.bar(labels, values)
        plt.title(f'主题强度 - {new_title}')
        plt.xlabel('主题')
        plt.ylabel('强度')
        plt.savefig(f'png/{new_title}.png')
        plt.show()


        # 生成csv文件
        df_strength = pd.DataFrame(
            {'主题': ['经济', '民生', '教育','社会','军事','文化','科技','政治'],
             '强度': [economic_strength, livelihood_strength, education_strength,society_strength,military_strength , culture_strength ,technology_strength,  politics_strength]})
        df_strength.to_csv(f'csv/{new_title}.csv', index=False)


